# contig_data.py

class ContigData:
    def __init__(self):
        self.data = {}
        # Initialize the 'No_Contig' category right away
        self.data['No_Contig'] = {
            'No_CS': 0  # For reads not aligning to known contigs and without a 'cs' tag
        }
    

    def add_contig(self, contig_name):
        # Initialize dictionaries for the new contig
        if contig_name not in self.data:
            self.data[contig_name] = {
                'substitutions': {},
                'deletions': {},
                'insertions': 0,
                'ambiguous_deletions': 0,
                'WT': 0,
                'No_CS': 0,  # Adding here the initialization for 'No_CS'
                'discarded': 0  # Adding the counter for discarded reads
            }
            
    def update_unaligned_read(self):
        # A simple counter for reads not aligned to known contigs
        self.data['No_Contig']['No_CS'] += 1

    def initialize_contigs(self, contig_ranges, consecutive_positions):
        for contig, (start, end, contig_length) in contig_ranges.items():  # Notice the change here
            self.add_contig(contig)
            
            # Prepare the set of positions to exclude
            exclude_positions = set(consecutive_positions.get(contig, []))

            for position in range(1, contig_length + 1):
                if position not in exclude_positions:
                    self.data[contig]['deletions'][position] = 0  # Initializes the count
                self.data[contig]['substitutions'][position] = 0  # Initializes the count regardless

    def update_substitution(self, contig_name, position, info):
        # Ensure the contig is initialized
        if contig_name not in self.data:
            self.add_contig(contig_name)

        substitutions = self.data[contig_name]['substitutions']
        if position in substitutions:
            # Increment the count of substitutions at this position
            substitutions[position] += 1  # incrementing instead of appending
        else:
            # Handle the case where the position is not initialized. This shouldn't happen if 'initialize_contigs' is used correctly.
            substitutions[position] = 1

    def update_deletion(self, contig_name, position):
        # Similar structure to `update_substitution`
        if contig_name not in self.data:
            self.add_contig(contig_name)

        deletions = self.data[contig_name]['deletions']
        if position in deletions:
            # Increment the count of deletions at this position
            deletions[position] += 1  # incrementing instead of appending
        else:
            # Handle the case where the position is not initialized. This shouldn't happen if 'initialize_contigs' is used correctly.
            deletions[position] = 1
        
    def update_simple_count(self, contig_name, mutation_type):
        # This method is updated to handle 'No_CS' within known contigs and 'No_Contig'
        if contig_name not in self.data:
            if contig_name == 'No_Contig':
                # We don't need to call add_contig for 'No_Contig', but we'll handle it here
                self.update_unaligned_read()
                return
            else:
                self.add_contig(contig_name)

        # For 'insertions', 'ambiguous_deletions', 'wt', and 'No_CS' (if within a known contig)
        self.data[contig_name][mutation_type] += 1

    def update_discarded(self, contig_name):
        # Increment the counter for discarded reads
        if contig_name not in self.data:
            self.add_contig(contig_name)

        self.data[contig_name]['discarded'] += 1

    def get_data(self):
        return self.data
